# PROBLEM 1

#

# This problem involves two tasks. 

# First, modify the acceleration function 

# to compute the acceleration of the 

# pendulum depending on its position.

# 

# Second, modify the symplectic_euler method 

# to determine the initial conditions and then 

# execute the Symplectic Euler Method.

#

# Please note that the order in which you 

# generate your initial conditions matters

# for grading purposes.  Have your first 

# (x, v) coordinate pair be the right-most

# point on the green ellipse, and progress

# counter-clockwise from there.

#




import numpy
import matplotlib.pyplot
import math

h = 0.05 # s

g = 9.81 # m / s2

length = 1. # m



x0 =  2.

y0 =  0



a = .25

b = 2.



#R = 



def acceleration(position):

    return -1 * g / length * math.sin( position / length )


def symplectic_euler(): 

    axes_x = matplotlib.pyplot.subplot(311)

    axes_x.set_ylabel('x in m')

    axes_v = matplotlib.pyplot.subplot(312)

    axes_v.set_ylabel('v in m/s')

    axes_v.set_xlabel('t in s')

    axes_phase_space = matplotlib.pyplot.subplot(313)

    axes_phase_space.set_xlabel('x in m')

    axes_phase_space.set_ylabel('v in m/s')

    num_steps = 80

    x = numpy.zeros(num_steps + 1) # m around circumference

    v = numpy.zeros(num_steps + 1) # m / s

    colors = [(0, 'g'), (3, 'c'), (10, 'b'), (30, 'm'), (79, 'r')]

    times = h * numpy.arange(num_steps + 1)



    num_initial_conditions = 50



    for i in range(num_initial_conditions):

        # Your code here

        x[0] = x0 + math.cos( i*2*math.pi/num_initial_conditions ) * a

        v[0] = y0 + math.sin( i*2*math.pi/num_initial_conditions ) * b

        

        for n in range( 1,num_steps ):

            x[n] = x[n-1]+h*v[n-1]

            v[n] = v[n-1]+h*acceleration(x[n])



        # Don't worry about this part of the function. It's just for making 

        # the plot look a bit nicer.

        axes_x.plot(times, x, c = 'k', alpha = 0.1)

        axes_v.plot(times, v, c = 'k', alpha = 0.1)        

        for step, color in colors:

            matplotlib.pyplot.hold(True)

            axes_x.scatter(times[step], x[step], c = color, edgecolors = 'none')

            axes_v.scatter(times[step], v[step], c = color, edgecolors = 'none')        

            axes_phase_space.scatter(x[step], v[step], c = color, edgecolors = 'none', s = 4)

    matplotlib.pyplot.show()


    return x, v



symplectic_euler()




